#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：domain-drop 
@File    ：res_draw.py
@IDE     ：PyCharm 
@Author  ：cao xu
@Date    ：2025/9/16 下午1:34 
"""
import re
import matplotlib.pyplot as plt

def parse_baseline_log(file_path):
    """解析 baseline 格式日志"""
    epochs, train_loss, val_loss, train_acc, val_acc = [], [], [], [], []
    pattern = re.compile(
        r"Epoch\s+(\d+)\s*\|\s*train_loss=([\d.]+)\s*acc=([\d.]+)\s*\|\s*val_loss=([\d.]+)\s*acc=([\d.]+)"
    )
    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            match = pattern.search(line)
            if match:
                epochs.append(int(match.group(1)))
                train_loss.append(float(match.group(2)))
                train_acc.append(float(match.group(3)))
                val_loss.append(float(match.group(4)))
                val_acc.append(float(match.group(5)))
    return epochs, train_loss, val_loss, train_acc, val_acc


def parse_step_log(file_path):
    """解析 step 级别日志 (20250905-train.txt)"""
    steps, loss, acc = [], [], []
    pattern = re.compile(
        r"(\d+)/\d+\s+of\s+epoch\s+(\d+)/\d+\s+class\s*:\s*([\d.]+),\s*domain\s*:\s*([\d.]+),\s*loss\s*:\s*([\d.]+)\s*-\s*acc class\s*:\s*([\d.]+)"
    )
    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            match = pattern.search(line)
            if match:
                step = int(match.group(1))
                epoch = int(match.group(2))
                loss_val = float(match.group(5))
                acc_val = float(match.group(6))
                steps.append(epoch * 10000 + step)  # 用 epoch*大数 + step 作为横坐标
                loss.append(loss_val)
                acc.append(acc_val)
    return steps, loss, acc


def plot_baseline(file_path, title_prefix):
    epochs, train_loss, val_loss, train_acc, val_acc = parse_baseline_log(file_path)

    # Loss 曲线
    plt.figure(figsize=(8,6))
    plt.plot(epochs, train_loss, label="Train Loss")
    plt.plot(epochs, val_loss, label="Val Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title(f"{title_prefix} - Loss Curve")
    plt.legend()
    plt.grid(True)
    plt.savefig(f"{title_prefix}_loss.png", dpi=300)
    plt.close()

    # Accuracy 曲线
    plt.figure(figsize=(8,6))
    plt.plot(epochs, train_acc, label="Train Acc")
    plt.plot(epochs, val_acc, label="Val Acc")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.title(f"{title_prefix} - Accuracy Curve")
    plt.legend()
    plt.grid(True)
    plt.savefig(f"{title_prefix}_acc.png", dpi=300)
    plt.close()


def plot_step(file_path, title_prefix):
    steps, loss, acc = parse_step_log(file_path)

    # Loss 曲线
    plt.figure(figsize=(8,6))
    plt.plot(steps, loss, label="Loss")
    plt.xlabel("Step")
    plt.ylabel("Loss")
    plt.title(f"{title_prefix} - Loss Curve")
    plt.legend()
    plt.grid(True)
    plt.savefig(f"{title_prefix}_loss.png", dpi=300)
    plt.close()

    # Accuracy 曲线
    plt.figure(figsize=(8,6))
    plt.plot(steps, acc, label="Acc")
    plt.xlabel("Step")
    plt.ylabel("Accuracy (%)")
    plt.title(f"{title_prefix} - Accuracy Curve")
    plt.legend()
    plt.grid(True)
    plt.savefig(f"{title_prefix}_acc.png", dpi=300)
    plt.close()


if __name__ == '__main__':
    plot_step("20250905-train.txt", "DomainDrop ResNet50")
    plot_baseline("train_baseline_thyroid_resnet50.txt", "Baseline ResNet50")
